"""
This simple example shows how you could use MLFlow REST API to create new 
runs inside an experiment to log parameters/metrics.  Using MLFlow REST API 
instead of MLFlow library might be useful to embed in an application where 
you don't want to depend on the whole mlflow library, or to make 
your own HTTP requests in another programming language (not Python).
For more details on MLFLow REST API endpoints check the following page: 

https://www.mlflow.org/docs/latest/rest-api.html 
"""

import argparse
import os
import time
import requests

_DEFAULT_USER_ID = "unknown"


class MLFlowTrackingRestApi:
	def __init__(self, hostname, port, experiment_id):
		self.base_url = 'http://' + hostname + ':' + str(port) + '/api/2.0/preview/mlflow'
		self.experiment_id = experiment_id
		self.run_id = self.create_run()

	def create_run(self):
		"""Create a new run for tracking."""
		url = self.base_url + '/runs/create'
		payload = {'experiment_id': self.experiment_id, 'start_time': int(time.time() * 1000), 'user_id': _get_user_id()}
		r = requests.post(url, json=payload)
		run_id = None
		if r.status_code == 200:
			run_id = r.json()['run']['info']['run_uuid']
		else:
			print("Creating run failed!")
		return run_id

	def list_experiments(self):
		"""Get all experiments."""
		url = self.base_url + '/experiments/list'
		r = requests.get(url)
		experiments = None
		if r.status_code == 200:
			experiments = r.json()['experiments']
		return experiments

	def log_param(self, param):
		"""Log a parameter dict for the given run."""
		url = self.base_url + '/runs/log-parameter'
		payload = {'run_uuid': self.run_id, 'key': param['key'], 'value': param['value']}
		r = requests.post(url, json=payload)
		return r.status_code

	def log_metric(self, metric):
		"""Log a metric dict for the given run."""
		url = self.base_url + '/runs/log-metric'
		payload = {'run_uuid': self.run_id, 'key': metric['key'], 'value': metric['value']}
		r = requests.post(url, json=payload)
		return r.status_code

def _get_user_id():
	"""Get the ID of the user for the current run."""
	try:
		import pwd
		return pwd.getpwuid(os.getuid())[0]
	except ImportError:
		return _DEFAULT_USER_ID


if __name__ == "__main__":
	# Command-line arguments
	parser = argparse.ArgumentParser(description='MLFlow REST API Example')

	parser.add_argument('--hostname', type=str, default='localhost', dest='hostname', 
		help='MLFlow server hostname/ip (default: localhost)')

	parser.add_argument('--port', type=int, default=5000, dest='port', 
		help='MLFlow server port number (default: 5000)')

	parser.add_argument('--experiment-id', type=int, default=0, dest='experiment_id', 
		help='Experiment ID (default: 0)')

	print("Running mlflow_tracking_rest_api.py")

	args = parser.parse_args()

	mlflow_rest = MLFlowTrackingRestApi(args.hostname, args.port, args.experiment_id)
	# Parameter is a key/val pair (str types)
	param = {'key': 'alpha', 'value': '0.1980'}
	status_code = mlflow_rest.log_param(param)
	if status_code == 200:
		print("Successfully logged parameter: {} with value: {}".format(param['key'], param['value']))
	else:
		print("Logging parameter failed!")
	# Metric is a key/val pair (key/val have str/float types)
	metric = {'key': 'precision', 'value': 0.769}	
	status_code = mlflow_rest.log_metric(metric)
	if status_code == 200:
		print("Successfully logged parameter: {} with value: {}".format(metric['key'], metric['value']))
	else:
		print("Logging metric failed!")
